#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-08-15


import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args, sms2pcl
from spirems.mod_helper import download_model
import argparse
import platform
from spirecv.algorithm.utils import calc_fov, calc_los_pos
from mmdet3d.apis import LidarDet3DInferencer
import copy
from scipy.spatial.transform import Rotation
from mmdet3d.structures import limit_period


def trans_det_results(det_results, camera_matrix, calib_wh, objs_in_meter):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = int(det_results.orig_shape[0])
    sms_results["width"] = int(det_results.orig_shape[1])
    has_calib = False
    if calib_wh[0] > 0 and calib_wh[1] > 0:
        sms_results["fov_x"], sms_results["fov_y"] = calc_fov(camera_matrix, calib_wh)
        has_calib = True

    sms_results["targets"] = []
    cls = det_results.boxes.cls.cpu().numpy().astype(np.float64)
    conf = det_results.boxes.conf.cpu().numpy().astype(np.float64)
    xywh = det_results.boxes.xywh.cpu().numpy().astype(np.float64)
    for i in range(len(cls)):
        ann = dict()
        name = det_results.names[int(cls[i])].strip()
        ann["category_name"] = name.replace(' ', '_').lower()
        ann["category_id"] = int(cls[i])
        ann["score"] = round(conf[i], 3)
        ann["bbox"] = [round(j, 3) for j in xywh[i].tolist()]
        ann["bbox"][0] = ann["bbox"][0] - ann["bbox"][2] / 2
        ann["bbox"][1] = ann["bbox"][1] - ann["bbox"][3] / 2
        ann["cxy"] = [
            (ann["bbox"][0] + ann["bbox"][2] / 2.) / sms_results["width"], 
            (ann["bbox"][1] + ann["bbox"][3] / 2.) / sms_results["height"]
        ]
        if has_calib and name in objs_in_meter:
            ann["los"], ann["pos"] = calc_los_pos(
                camera_matrix, calib_wh, 
                ann["cxy"], [ann["bbox"][2], ann["bbox"][3]], 
                objs_in_meter[name]
            )
        sms_results["targets"].append(ann)

    return sms_results


class PointPillars3DDetNode_Cuda(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.launch_next_emit = self.get_param("launch_next_emit", True)
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.model = self.get_param("model", "sms::pointpillars_hv_secfpn_8xb6-160e_kitti-3d-car.py")
        self.weights = self.get_param("weights", "sms::hv_pointpillars_secfpn_6x8_160e_kitti-3d-car_20220331.pth")
        self.device = self.get_param("device", "cuda:0")
        self.confidence = self.get_param("confidence", 0.3)
        self.dataset_name = self.get_param("dataset_name", "kitti")  # kitti, nuscenes_c10
        self.class_filter = self.get_param("class_filter", [], "设置数据集中需要检测的类别，空列表则表示检测所有类别")
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.names = self.dataset_categories[self.dataset_name]

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True


        if self.model.startswith("sms::"):
            self.local_model = download_model(self.__class__.__name__, self.model)
            assert self.local_model is not None
        else:
            self.local_model = self.model

        if self.weights.startswith("sms::"):
            self.local_weights = download_model(self.__class__.__name__, self.weights)
            assert self.local_weights is not None
        else:
            self.local_weights = self.weights


        input_url = '/' + job_name + '/sensor/pointcloud'
        if len(self.specified_input_topic) > 0:
            input_url = self.specified_input_topic

        output_url = '/' + job_name + '/detector/results'
        if len(self.specified_output_topic) > 0:
            output_url = self.specified_output_topic

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self._pcd_reader = Subscriber(
            input_url, 'memory_msgs::PointCloud', self.pcd_callback,
            ip=ip, port=port
        )
        self._result_writer = Publisher(
            output_url, 'geometry_msgs::SceneUpdate',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )
        
        self.inferencer = LidarDet3DInferencer(
            model=self.local_model,
            weights=self.local_weights,
            device=self.device
        )
        self.start()

    def release(self):
        BaseNode.release(self)
        self._pcd_reader.kill()
        self._result_writer.kill()
        self._next_writer.kill()

    def pcd_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        pcl = sms2pcl(msg)
        self.job_queue.put({'msg': msg, 'pcl': pcl})

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            msg, pcl = msg_dict['msg'], msg_dict['pcl']
            file_name = msg['file_name'] if 'file_name' in msg else ''
            print(pcl.shape)

            results_dict = self.inferencer(
                inputs={"points": pcl},
                pred_score_thr=self.confidence
            )["predictions"][0]
            # print(results_dict)

            entity = def_msg("geometry_msgs::SceneEntity")
            entity["frame_id"] = msg["frame_id"]
            entity["frame_locked"] = True
            entity["lifetime"]["sec"] = 1
            entity["cubes"] = []
            for label, score, bbox_3d in zip(results_dict["labels_3d"], results_dict["scores_3d"], results_dict["bboxes_3d"]):
                if score > self.confidence:
                    if len(self.class_filter) == 0 or self.names[label] in self.class_filter:
                        cube = def_msg("geometry_msgs::CubePrimitive")
                        cube["pose"]["position"]["x"] = bbox_3d[0]
                        cube["pose"]["position"]["y"] = bbox_3d[1]
                        cube["pose"]["position"]["z"] = bbox_3d[2] + bbox_3d[5] / 2
                        cube["pose"]["orientation"]["x"] = np.cos(bbox_3d[6] / 2)
                        cube["pose"]["orientation"]["y"] = np.sin(bbox_3d[6] / 2)
                        cube["pose"]["orientation"]["z"] = 0
                        cube["pose"]["orientation"]["w"] = 0
                        cube["size"]["x"] = bbox_3d[3]
                        cube["size"]["y"] = bbox_3d[4]
                        cube["size"]["z"] = bbox_3d[5]
                        cube["color"]["r"] = 255
                        cube["color"]["a"] = 0.2
                        entity["cubes"].append(copy.deepcopy(cube))

            vis_msg = def_msg("geometry_msgs::SceneUpdate")
            vis_msg["entities"] = [entity]
            self._result_writer.publish(vis_msg)

            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = PointPillars3DDetNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
